#! /usr/bin/env python3

from package.litedb import litedb
import matplotlib.pyplot as plt
import numpy as np
import sys
import time


def draw_memleak(task, pid=None):

    def format_coord(x, y):
        y = int(y + 0.5)
        xp = (x - xlim[0]) * 100 / (xlim[1] - xlim[0])
        return f'x = {x:.2f} ({xp:.2f}%)    y = {y}    '

    def on_scroll(event):
        ax = event.inaxes
        xmin, xmax = ax.get_xlim()
        print(f"{xmin}, {xmax}")

    def on_pick(event):
        ind = event.ind
        art = event.artist

        if not annot.get_visible():
            annot.set_visible(True)
        annot.xy = art.get_offsets()[ind[-1]]
        annot.set_text(f'time: {annot.xy[0]}\ndots: {len(ind)}')
        fig.canvas.draw_idle()
        print(annot.xy[0])

    db = litedb()
    db.open('trace.db')
    fig, ax = plt.subplots()

    xlim = []
    ylim = []

    res = db.sql(f"select time, bytes from 'trace' where task='{task}' and flag0='O'")
    arr = np.array(res.fetchall())
    if len(arr) != 0:
        xdat = arr[:, 0] # time
        ydat = arr[:, 1] # bytes
        xlim = [xdat[0], xdat[-1]]
        ylim = [min(ydat), max(ydat)]
        print(f'safe dots: {len(xdat)}')
        plt.scatter(xdat, np.log2(ydat), s=40, alpha=0.7, linewidths=0, picker=True, label='safe', c='cornflowerblue') # draw safe

    res = db.sql(f"select time, bytes from 'trace' where task='{task}' and flag0='X'")
    arr = np.array(res.fetchall())
    if len(arr) != 0:
        xdat = arr[:, 0]
        ydat = arr[:, 1]
        if not xlim:
            xlim = [xdat[0], xdat[-1]]
            ylim = [min(ydat), max(ydat)]
        else:
            xlim = [min(xlim[0], xdat[0]), max(xlim[1], xdat[-1])]
            ylim = [min(ylim[0], min(ydat)), max(ylim[1], max(ydat))]
        print(f'leak dots: {len(xdat)}')
        plt.scatter(xdat, np.log2(ydat)+0.2, s=40, alpha=0.7, linewidths=0, picker=True, label='leak', c='darkorange') # draw leak

    # set axis
    xl = xlim
    xd = (xl[1]-xl[0])*0.1
    xl = [xl[0]-xd, xl[1]+xd]
    yl = np.log2(ylim)
    yd = 1
    yl = [yl[0]-yd, yl[1]+yd]
    ax.set(xlim=xl, ylim=yl, title=f'{task}', xlabel='Time', ylabel='Slab Size')
    ax.ticklabel_format(useOffset=False) #, style='plain')
    ax.format_coord = format_coord

    # init annotate
    annot = ax.annotate("", xy=(0,0), xytext=(-10,30), textcoords='offset points', 
            bbox=dict(boxstyle='square', fc='w'), arrowprops=dict(arrowstyle='->'))
    annot.set_visible(False)

    # set figure
    fig.subplots_adjust(top=0.9, bottom=0.2) # bottom blank
    fig.canvas.mpl_connect('pick_event', on_pick)
    #fig.canvas.mpl_connect('scroll_event', on_scroll)
    #fig.canvas.mpl_connect('motion_notify_event', on_hover)

    # set plot and draw
    plt.xticks(rotation=15)
    plt.legend()
    plt.grid(True)
    plt.show()


if __name__=='__main__':
    db = litedb()
    db.open('trace.db')
    #draw_memleak_by_task(sys.argv[1])
    draw_memleak(sys.argv[1])
    db.close()
